# -*- coding: utf-8 -*-
# author： woldier wong
# datetime： 3/12/24 3:15 AM
# ide： PyCharm

import torch
import math
from torch import nn, Tensor
from typing import Any, Dict, List, Optional, Tuple
from torchvision.models.detection.image_list import ImageList
import torchvision
from torchvision.models.detection.roi_heads import paste_masks_in_image  # TODO 改为自己的code


@torch.jit.unused
def _get_shape_onnx(image: Tensor) -> Tensor:
    from torch.onnx import operators

    return operators.shape_as_tensor(image)[-2:]


@torch.jit.unused
def _fake_cast_onnx(v: Tensor) -> float:
    # ONNX requires a tensor but here we fake its type for JIT.
    return v


def _resize_image_and_masks(
        image: Tensor,
        self_min_size: float,
        self_max_size: float,
        target: Optional[Dict[str, Tensor]] = None,
        fixed_size: Optional[Tuple[int, int]] = None,
) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
    """
    缩放图像和掩码图
    :param image:
    :param self_min_size:
    :param self_max_size:
    :param target:
    :param fixed_size:
    :return:
    """
    if torchvision._is_tracing():
        im_shape = _get_shape_onnx(image)
    else:
        im_shape = torch.tensor(image.shape[-2:])

    size: Optional[List[int]] = None
    scale_factor: Optional[float] = None  # 缩放因子
    recompute_scale_factor: Optional[bool] = None
    if fixed_size is not None:
        size = [fixed_size[1], fixed_size[0]]
    else:
        min_size = torch.min(im_shape).to(dtype=torch.float32)  # image 高宽中较小的一个
        max_size = torch.max(im_shape).to(dtype=torch.float32)  # image 高宽中较大的一个
        scale = torch.min(self_min_size / min_size, self_max_size / max_size)  # 计算合适的缩放系数

        if torchvision._is_tracing():
            scale_factor = _fake_cast_onnx(scale)
        else:
            scale_factor = scale.item()  # 得到缩放因子
        recompute_scale_factor = True
        # image resize
    image = torch.nn.functional.interpolate(
        image[None],
        size=size,  # 输出大小
        scale_factor=scale_factor,  # 缩放因子
        mode="bilinear",
        recompute_scale_factor=recompute_scale_factor,  # 是否让缩放因子占主导地位, 而不是size
        align_corners=False,
    )[0]

    if target is None:
        return image, target

    if "masks" in target:
        mask = target["masks"]
        mask = torch.nn.functional.interpolate(
            mask[:, None].float(), size=size, scale_factor=scale_factor, recompute_scale_factor=recompute_scale_factor
        )[:, 0].byte()
        target["masks"] = mask
    return image, target


class GeneralizedRCNNTransform(nn.Module):
    """
    Performs input / target transformation before feeding the data to a GeneralizedRCNN
    model.
    在将数据送入GeneralizedRCNN之前, 对 input/target进行Transform处理
    The transformations it perform are:
        - input normalization (mean subtraction and std division)  input 图像的norm
        - input / target resizing to match min_size / max_size  input 图像和target 的resize, 这里的resize 是等比例缩放
                需要注意的是, targets中 box 的key 为 boxes
    It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets
    """

    def __init__(
            self,
            min_size: int,
            max_size: int,
            image_mean: List[float],
            image_std: List[float],
            size_divisible: int = 32,
            fixed_size: Optional[Tuple[int, int]] = None,
            **kwargs: Any,
    ):
        """

        :param min_size: resize 图片的最小长宽
        :param max_size:  resize 图片的最大长宽
        :param image_mean: 图像的均值 (m1,m2,m3)
        :param image_std:  图像的标准差 (s1,s2,s3)
        :param size_divisible:  resize后图片的 整体h w 必须可以被  size_divisible整除
        :param fixed_size:
        :param kwargs:
        """
        super().__init__()
        if not isinstance(min_size, (list, tuple)):  # 如果min_size不是一个list 或者是tuple 那么将其转换为tuple
            min_size = (min_size,)
        self.min_size = min_size
        self.max_size = max_size
        self.image_mean = image_mean
        self.image_std = image_std
        self.size_divisible = size_divisible
        self.fixed_size = fixed_size
        self._skip_resize = kwargs.pop("_skip_resize", False)  # 弹出一个kwargs, 默认为False

    def forward(
            self, images: List[Tensor], targets: Optional[List[Dict[str, Tensor]]] = None
    ) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]:
        """
        正向传播过程
        :param images: 输入的图像, 由于batch中图像的尺寸可能不一样, 因此不能打包成一个大的tensor而是用一个list数组来装的.
            images[i]的shape 是[C_i,H_i,W_i]
        :param targets: GT box, 也是用一个list来存的不同图片的GT box 信息
            例如 [ {'classes':Tensor([0,2]), 'bboxes':Tensor([ [12,32,56,88], [14,25,77,98] ] )} ,]
        :return:
        """
        # 遍历images 数组
        images = [img for img in images]
        if targets is not None:  # 如果targets非None 那么说明是训练模式
            # make a copy of targets to avoid modifying it in-place
            # once torchscript supports dict comprehension
            # this can be simplified as follows
            # targets = [{k: v for k,v in t.items()} for t in targets]
            targets_copy: List[Dict[str, Tensor]] = []
            for t in targets:
                data: Dict[str, Tensor] = {}  # 存放一张图片中目标信息的dict
                for k, v in t.items():  # 遍历t.items() 即kv 对
                    data[k] = v
                targets_copy.append(data)  # 放入list
            targets = targets_copy
        for i in range(len(images)):  # 对batch 中的每一张图片进行处理
            image = images[i]  # 获取第i张图片
            target_index = targets[i] if targets is not None else None  # 获取第i个target , 注意在inference时有可能是None

            if image.dim() != 3:  # 判断image的shape 是否合理
                raise ValueError(f"images is expected to be a list of 3d tensors of shape [C, H, W], got {image.shape}")
            image = self.normalize(image)  # 对image 进行归一化
            image, target_index = self.resize(image, target_index)  # 进行缩放, 需要同时处理图片以及对应的GT box 中 mask 信息
            images[i] = image  # 重新设置处理后的image
            if targets is not None and target_index is not None:
                targets[i] = target_index

        image_sizes = [img.shape[-2:] for img in images]  # 得到resize 后的image 的h,w 信息
        # 将list中的不同size 的image 进行padding 形成一个batch tensor [B, C, H, W]
        images = self.batch_images(images, size_divisible=self.size_divisible)
        image_sizes_list: List[Tuple[int, int]] = []  # 记录resize 之后, padding 之前的image size
        for image_size in image_sizes:
            torch._assert(
                len(image_size) == 2,
                f"Input tensors expected to have in the last two elements H and W, instead got {image_size}",
            )
            image_sizes_list.append((image_size[0], image_size[1]))

        image_list = ImageList(images, image_sizes_list)
        return image_list, targets

    def normalize(self, image: Tensor) -> Tensor:
        """

        对一张image进行归一化
        :param image: shape [C, H, W]
        :return:
        """
        if not image.is_floating_point():  # 判断是否是0-1之间的数值
            raise TypeError(
                f"Expected input images to be of floating type (in range [0, 1]), "
                f"but found type {image.dtype} instead"
            )
        dtype, device = image.dtype, image.device  # 得到image 的dtype 和 device

        # 创建与image 属性相同的tensor
        mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)
        std = torch.as_tensor(self.image_std, dtype=dtype, device=device)
        # [3] -> [3, 1, 1] 然后通过广播机制与image 进行计算
        return (image - mean[:, None, None]) / std[:, None, None]

    def torch_choice(self, k: List[int]) -> int:
        """
        Implements `random.choice` via torch ops so it can be compiled with
        TorchScript. Remove if https://github.com/pytorch/pytorch/issues/25803
        is fixed.
        """
        index = int(torch.empty(1).uniform_(0.0, float(len(k))).item())
        return k[index]

    def resize(
            self,
            image: Tensor,
            target: Optional[Dict[str, Tensor]] = None,
    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
        """
        将图片进行等比例缩放
        :param image: 图像 shape is [C, H, W]
        :param target: GT box 信息  {'classes':Tensor([0,2]), 'bboxes':Tensor([ [12,32,56,88], [14,25,77,98] ] )}
        :return:
        """
        h, w = image.shape[-2:]  # 获取图片的宽高
        if self.training:  # 如果是train 模式
            if self._skip_resize:  # 如果是跳过resize
                return image, target
            size = float(self.torch_choice(self.min_size))
        else:
            # FIXME assume for now that testing uses the largest scale
            size = float(self.min_size[-1])
        image, target = _resize_image_and_masks(image, size, float(self.max_size), target, self.fixed_size)

        if target is None:
            return image, target
        # 对box进行操作
        bbox = target["boxes"]
        bbox = resize_boxes(bbox, (h, w), image.shape[-2:])
        target["boxes"] = bbox

        if "keypoints" in target:
            keypoints = target["keypoints"]
            keypoints = resize_keypoints(keypoints, (h, w), image.shape[-2:])
            target["keypoints"] = keypoints
        return image, target

    # _onnx_batch_images() is an implementation of
    # batch_images() that is supported by ONNX tracing.
    @torch.jit.unused
    def _onnx_batch_images(self, images: List[Tensor], size_divisible: int = 32) -> Tensor:
        max_size = []
        for i in range(images[0].dim()):
            max_size_i = torch.max(torch.stack([img.shape[i] for img in images]).to(torch.float32)).to(torch.int64)
            max_size.append(max_size_i)
        stride = size_divisible
        max_size[1] = (torch.ceil((max_size[1].to(torch.float32)) / stride) * stride).to(torch.int64)
        max_size[2] = (torch.ceil((max_size[2].to(torch.float32)) / stride) * stride).to(torch.int64)
        max_size = tuple(max_size)

        # work around for
        # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
        # which is not yet supported in onnx
        padded_imgs = []
        for img in images:
            padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
            padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
            padded_imgs.append(padded_img)

        return torch.stack(padded_imgs)

    def max_by_axis(self, the_list: List[List[int]]) -> List[int]:
        """
        统计list中 H,W 的最大值
        :param the_list:
        :return:
        """
        maxes = the_list[0]  # 设置当前的最大值为0
        for sublist in the_list[1:]:
            for index, item in enumerate(sublist):  # 统计H , W 的最大值
                maxes[index] = max(maxes[index], item)
        return maxes

    def batch_images(self, images: List[Tensor], size_divisible: int = 32) -> Tensor:
        """
        将一个list 的image cat成一个大的tensor
        :param images:
        :param size_divisible:
        :return:
        """
        if torchvision._is_tracing():
            # batch_images() does not export well to ONNX
            # call _onnx_batch_images() instead
            return self._onnx_batch_images(images, size_divisible)
        # 生成一个保存有image c, h, w的list 输入函数 得到C, H, W的最大值
        max_size = self.max_by_axis([list(img.shape) for img in images])
        stride = float(size_divisible)  # 步长
        max_size = list(max_size)  # 一维数组长度为3, 里面的值为 [C_max,H_max,W_max]
        max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride)  # 向上取整到可以被size_divisible整除的最小整数
        max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride)

        batch_shape = [len(images)] + max_size  # 数组拼接成一个长度为4的一维数组 [batch_size,C_max,H_max',W_max']
        batched_imgs = images[0].new_full(batch_shape, 0)  # 填充一个 shape 为 [batch_size,C_max,H_max',W_max'] 值全为0的tensor
        for i in range(batched_imgs.shape[0]):  # 将image与背景(batched_imgs)进行组合
            img = images[i]
            batched_imgs[i, : img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)  # 拷贝

        return batched_imgs

    def postprocess(
            self,
            result: List[Dict[str, Tensor]],
            image_shapes: List[Tuple[int, int]],
            original_image_sizes: List[Tuple[int, int]],
    ) -> List[Dict[str, Tensor]]:
        """
        对于变换的逆变换
        :param result: 推理结果, 包括boxes, keypoint, mask
        :param image_shapes:
        :param original_image_sizes:
        :return:
        """
        # 如果是train, 则不用进行post transform
        if self.training:
            return result
        for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):
            boxes = pred["boxes"]
            boxes = resize_boxes(boxes, im_s, o_im_s)  # 将boxes 转换为原始的尺度
            result[i]["boxes"] = boxes
            if "masks" in pred:
                masks = pred["masks"]
                masks = paste_masks_in_image(masks, boxes, o_im_s)
                result[i]["masks"] = masks
            if "keypoints" in pred:
                keypoints = pred["keypoints"]
                keypoints = resize_keypoints(keypoints, im_s, o_im_s)
                result[i]["keypoints"] = keypoints
        return result

    def __repr__(self) -> str:
        format_string = f"{self.__class__.__name__}("
        _indent = "\n    "
        format_string += f"{_indent}Normalize(mean={self.image_mean}, std={self.image_std})"
        format_string += f"{_indent}Resize(min_size={self.min_size}, max_size={self.max_size}, mode='bilinear')"
        format_string += "\n)"
        return format_string


def resize_boxes(boxes: Tensor, original_size: List[int], new_size: List[int]) -> Tensor:
    """
    经过了缩放之后, resize box的大小
    :param boxes: box值
    :param original_size: 原始的图片大小
    :param new_size: 经过resize 后的图片大小
    :return:
    """
    ratios = [
        torch.tensor(s, dtype=torch.float32, device=boxes.device)
        / torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)
        for s, s_orig in zip(new_size, original_size)
    ]
    ratio_height, ratio_width = ratios
    xmin, ymin, xmax, ymax = boxes.unbind(1)

    xmin = xmin * ratio_width
    xmax = xmax * ratio_width
    ymin = ymin * ratio_height
    ymax = ymax * ratio_height
    return torch.stack((xmin, ymin, xmax, ymax), dim=1)


def resize_keypoints(keypoints: Tensor, original_size: List[int], new_size: List[int]) -> Tensor:
    ratios = [
        torch.tensor(s, dtype=torch.float32, device=keypoints.device)
        / torch.tensor(s_orig, dtype=torch.float32, device=keypoints.device)
        for s, s_orig in zip(new_size, original_size)
    ]
    ratio_h, ratio_w = ratios
    resized_data = keypoints.clone()
    if torch._C._get_tracing_state():
        resized_data_0 = resized_data[:, :, 0] * ratio_w
        resized_data_1 = resized_data[:, :, 1] * ratio_h
        resized_data = torch.stack((resized_data_0, resized_data_1, resized_data[:, :, 2]), dim=2)
    else:
        resized_data[..., 0] *= ratio_w
        resized_data[..., 1] *= ratio_h
    return resized_data


if __name__ == "__main__":
    trans = GeneralizedRCNNTransform(
        min_size=256,
        max_size=512,
        image_mean=[0.231, 0.312, 0.112],
        image_std=[0.231, 0.312, 0.112],
    )
    images = [torch.randn((3, h, w)) for h, w in zip((680, 430, 512), (415, 513, 512))]
    targets = [
        {'classes': Tensor([0, 2]), 'boxes': Tensor([[12, 32, 56, 88], [14, 25, 77, 98]])},
        {'classes': Tensor([0, 2]), 'boxes': Tensor([[12, 32, 56, 88], [14, 25, 77, 98]])},
        {'classes': Tensor([0, 2]), 'boxes': Tensor([[12, 32, 56, 88], [14, 25, 77, 98]])},
    ]

    images, targets = trans(images, targets)


